# 创建时间：2024/7/9下午1:21
from utils import *


class Resolver:
    """
    包解析器
    实现功能：
        十六进制转储
        pcap包解析（目前只支持解析单个包头，解析多包头待完善）
        IP协议解析（只支持IPv4）
        TCP协议解析
    """

    def __init__(self, data):
        """
        解析器构造

        :param data: 要解析的数据，可以是字节串、字节流、字节列表、字符串，传入后以字节串保存
        """
        # 设置初始状态
        self.pcap = None
        """如果数据可以以pcap格式解析，保存pcap信息，否则为None"""
        self.ip = None
        """如果数据可以以ip格式解析，保存ip信息，否则为None"""
        self.tcp = None
        """如果数据可以以tcp格式解析，保存tcp信息，否则为None"""

        import _io
        # 如果接收到字节串，则保存
        if isinstance(data, bytes):
            self.data = data
        # 如果接收到字符串，则对字符串编码
        elif isinstance(data, str):
            self.data = data.encode()
        # 如果接收到字节流，则读取所有行
        elif isinstance(data, _io.BufferedReader):
            data = data.readlines()
        # 如果接收到列表，则遍历列表内容存入字节串
        if isinstance(data, list):
            cache = bytes()
            for i in data:
                cache += i
            self.data = cache

    class _PCAP(object):
        """
        用于保存解析到的pcap信息
        """

        def __new__(cls, *args, **kwargs):
            # 如果数据总长度小于pcap首部长度则返回空
            if len(args[0]) <= 40:
                return None
            return super().__new__(cls)

        def __init__(self, byte_str: bytes):
            self.data = byte_str
            """原始字节串"""
            self.length = len(byte_str)
            """字节串总长度"""
            self.little_order = False
            """字节序是否小端序"""
            self.nanoseconds = False
            """时间戳精度是否纳秒"""
            # 字节串转字节列表，4byte一行拆分字节串
            byte_list: list[bytes | list] = [byte_str[i * 4:(i + 1) * 4] for i in range(0, 10)]
            byte_list.append(byte_str[40:])
            # 拆分主次版本号
            byte_list[1] = [byte_list[1][:2], byte_list[1][2:]]
            # 第一行保存到幻数
            magic_number = byte_list[0].hex()
            # 根据幻数判断字节序和时间戳精度
            match magic_number:
                case 'd4c3b2a1':
                    self.little_order = True
                    self.nanoseconds = False
                case '4d3cb2a1':
                    self.little_order = True
                    self.nanoseconds = True
                case 'a1b2c3d4':
                    self.little_order = False
                    self.nanoseconds = False
                case 'a1b23c4d:':
                    self.little_order = False
                    self.nanoseconds = True
                case _:
                    raise ValueError('Magic Number字段值不匹配')
            # 如果是小端序，转大端序
            if self.little_order:
                for i in range(0, 10):
                    if i == 1:
                        for j in range(0, 2):
                            byte_list[i][j] = self.reverse_byte_order(byte_list[i][j])
                        continue
                    byte_list[i] = self.reverse_byte_order(byte_list[i])
            # 保存数据的字节列表
            self.data_list: list[bytes | list] = byte_list
            """按4byte拆分的字节列表"""
            # 保存幻数
            self.magic_number: str = byte_list[0].hex()
            """幻数"""
            # 保存主版本号
            self.major_version: str = str(int.from_bytes(byte_list[1][0], 'big', signed=False))
            """主版本号"""
            # 保存次版本号
            self.minor_version: str = str(int.from_bytes(byte_list[1][1], 'big', signed=False))
            """次版本号"""
            # 保存接口快照长度
            self.snap_len: str = str(int.from_bytes(byte_list[4], 'big', signed=False))
            """接口快照长度/截断长度"""
            # 保存链路类型
            self.link_type = int.from_bytes(byte_list[5], 'big', signed=False)
            """链路类型"""
            # 解析时间戳
            import datetime
            t = int.from_bytes(byte_list[6], 'big', signed=False) + \
                (int.from_bytes(byte_list[7], 'big', signed=False) / 1000000000
                 if self.nanoseconds else
                 int.from_bytes(byte_list[7], 'big', signed=False) / 1000000)
            self.timestamp = datetime.datetime.fromtimestamp(t)
            """时间戳"""
            # 保存数据包捕获长度
            self.captured_packet_length: str = str(int.from_bytes(byte_list[8], 'big', signed=False))
            """数据包捕获长度"""
            # 保存数据包原始长度
            self.original_packet_length: str = str(int.from_bytes(byte_list[9], 'big', signed=False))
            """数据包原始长度"""
            # 保存包数据
            self.packet_data: bytes = self.data_list[10]
            """包数据"""

        def __repr__(self):
            result = 'PCAP'.center(40, '=') + '\n'
            result += '字节长度 : ' + str(self.length) + '\n'
            result += '全局报头'.center(30, '-') + '\n'
            result += '主版本号 : ' + self.major_version + '\n'
            result += '次版本号 : ' + self.minor_version + '\n'
            result += '截断长度 : ' + self.snap_len + '\n'
            result += '链路类型 : ' + str(self.link_type) + '\n'
            result += '数据记录'.center(30, '-') + '\n'
            result += '时间戳 : ' + str(self.timestamp) + '\n'
            result += '捕获长度 : ' + self.captured_packet_length + '\n'
            result += '原始长度 : ' + self.original_packet_length + '\n'
            return result

        def get_report(self, showdata: bool = False):
            """
            打印pcap解析结果
            """

            print('PCAP'.center(40, '='))
            print('字节长度\t: ' + str(self.length))
            print('全局报头'.center(30, '-'))
            print('主版本号\t: ' + self.major_version)
            print('次版本号\t: ' + self.minor_version)
            print('截断长度\t: ' + self.snap_len)
            print('链路类型\t: ' + str(self.link_type))
            print('数据记录'.center(30, '-'))
            print('时间戳\t: ' + str(self.timestamp))
            print('捕获长度\t: ' + self.captured_packet_length)
            print('原始长度\t: ' + self.original_packet_length)
            if showdata:
                print('数据'.center(30, '-'))
                hex_dump(self.packet_data)

        @staticmethod
        def reverse_byte_order(byte_str: bytes):
            """
            转换2byte、4byte或8byte字节串的字节序

            :param byte_str: 待转字节序的字节串，必须是2byte、4byte或8byte
            :return: 转换后的字节串
            """
            import struct
            if len(byte_str) == 2:
                unpacked_value = struct.unpack('<H', byte_str)[0]
                format_code = '>H'
            elif len(byte_str) == 4:
                unpacked_value = struct.unpack('<I', byte_str)[0]
                format_code = '>I'
            elif len(byte_str) == 8:
                unpacked_value = struct.unpack('<Q', byte_str)[0]
                format_code = '>Q'
            else:
                raise ValueError('只能转换2byte、4byte或8byte的字节串')
            return struct.pack(format_code, unpacked_value)

    class _IP(object):
        """用于保存解析到的IP协议信息"""

        def __new__(cls, *args, **kwargs):
            byte_str: bytes = args[0]
            # 如果数据长度小于IP数据报首部长度则返回空
            if len(byte_str) <= (byte_str[0] & 0xf) * 4:
                return None
            return super().__new__(cls)

        def __init__(self, byte_str: bytes):
            self.data = byte_str
            """原始字节串"""
            self.length = len(byte_str)
            """字节串总长度"""
            # 判断协议版本，不是IPv4则报错
            if (byte_str[0] & 0xf0) // 0xf != 4:
                raise ValueError('IP数据报只支持解析IPv4，不是IP数据报或版本不是IPv4')
            self.version = '4'
            """IP数据报版本"""
            self.ihl = byte_str[0] & 0xf
            """IP数据报首部长度，单位4byte"""
            # 计算可选字段行数
            self.count_options = self.ihl - 5
            """可选字段行数，每行4byte"""
            # 保存总长度
            self.total_length = int.from_bytes(byte_str[2:4])
            """IP数据报总长度"""
            # 如果字节串长度小于数据报总长度则报错
            if self.length < self.total_length:
                raise ValueError('字节串长度小于数据报总长度')
            # 字节串切分字节列表
            byte_list = [byte_str[i * 4:(i + 1) * 4] for i in range(5 + self.count_options)]
            byte_list.append(byte_str[(5 + self.count_options) * 4:])
            self.data_list = byte_list
            """字节列表，头部4byte一组"""
            self.identification = int.from_bytes(self.data_list[1][0:2], 'big', signed=False)
            """IP数据报标识"""
            self.sign_df = bool(int(format(self.data_list[1][2], '08b')[1]))
            """IP数据报DF标志位，是否禁止分片"""
            self.sign_mf = bool(int(format(self.data_list[1][2], '08b')[2]))
            """IP数据报MF标志位，是否有后续分片"""
            self.fragment_offset = Convert.to_int(Convert.to_str(self.data_list[1][2:4], 2)[3:], 2) * 8
            """IP数据报片偏移，单位1byte"""
            self.ttl = self.data_list[2][0]
            """IP数据报生存时间，剩余转发跳数"""
            self.protocol = ''
            """IP数据报协议"""
            protocol = self.data_list[2][1]
            # 解析协议
            match protocol:
                case 1:
                    self.protocol = 'ICMP'
                case 4:
                    self.protocol = 'IP'
                case 6:
                    self.protocol = 'TCP'
                case 17:
                    self.protocol = 'UDP'
                case _:
                    raise ValueError('协议类型错误，IP数据报解析失败')
            self.header_checksum = self.data_list[2][2:4]
            """IP数据报首部校验和，未做校验"""
            source_ip = self.data_list[3]
            self.source_ip = '.'.join([Convert.to_str(source_ip[i]) for i in range(4)])
            """IP数据报源IP地址"""
            destination_ip = self.data_list[4]
            self.destination_ip = '.'.join([Convert.to_str(destination_ip[i]) for i in range(4)])
            """IP数据报目的IP地址"""
            self.options = self.data_list[5:5 + self.count_options]
            """IP数据报可选字段+填充字段"""
            self.packet_data = self.data_list[5 + self.count_options]
            """IP数据报数据部分"""

        def get_report(self, showdata=False):
            print('IP'.center(40, '='))
            print('字节长度\t: ' + str(self.length))
            print('首部'.center(30, '-'))
            print('版本\t\t: ' + self.version)
            print('首部长度\t: ' + str(self.ihl * 4))
            print('总长度\t: ' + str(self.total_length))
            print('标识\t\t: ' + str(self.identification))
            print('禁止分片\t: ' + ('是' if self.sign_df else '否'))
            print('后续分片\t: ' + ('有' if self.sign_mf else '无'))
            print('片偏移\t: ' + str(self.fragment_offset))
            print('剩余跳数\t: ' + str(self.ttl))
            print('协议\t\t: ' + self.protocol)
            print('首校验和\t: ' + Convert.to_str(self.header_checksum, 16))
            print('源地址\t: ' + self.source_ip)
            print('目的地址\t: ' + self.destination_ip)
            print('可选字段\t: ' + str(self.options))
            if showdata:
                print('数据'.center(30, '-'))
                hex_dump(self.packet_data)

    class _TCP(object):
        """用于保存解析到的TCP协议信息"""

        def __new__(cls, *args, **kwargs):
            byte_str: bytes = args[0]
            # 如果数据长度小于IP数据报首部长度则返回空
            if len(byte_str) <= 20:
                return None
            return super().__new__(cls)

        def __init__(self, byte_str: bytes):
            self.data = byte_str
            """原始字节串"""
            self.length = len(byte_str)
            """字节串总长度"""
            self.data_offset = (byte_str[12] & 0xf0) // 0xf
            """TCP数据报数据偏移，单位4byte"""
            # 计算可选字段行数
            self.count_options = self.data_offset - 5
            """可选字段行数，每行4byte"""

            # 字节串切分字节列表
            byte_list = [byte_str[i * 4:(i + 1) * 4] for i in range(5 + self.count_options)]
            byte_list.append(byte_str[(5 + self.count_options) * 4:])
            self.data_list = byte_list
            """字节列表，头部4byte一组"""
            self.src_port = Convert.to_int(self.data_list[0][0:2])
            """TCP数据报源端口"""
            self.des_port = Convert.to_int(self.data_list[0][2:4])
            """TCP数据报目的端口"""
            self.seq_number = Convert.to_int(self.data_list[1])
            """TCP数据报序号"""
            self.ack_number = Convert.to_int(self.data_list[2])
            """TCP数据报确认号"""
            sign = Convert.to_str(self.data_list[3][1], 2, 8)[2:]
            # TCP数据报标志位赋值
            self.urg, self.ack, self.psh, self.pst, self.syn, self.fin = tuple([bool(int(i)) for i in sign])
            """TCP数据报标志位"""
            self.win_size = Convert.to_int(self.data_list[3][2:4])
            """TCP数据报窗口大小"""
            self.checksum = self.data_list[4][:2]
            """TCP数据报校验和，未做校验"""
            self.urg_pointer = Convert.to_int(self.data_list[4][2:])
            """TCP数据报紧急指针"""
            self.options = self.data_list[5:5 + self.count_options]
            """TCP数据报可选字段"""
            self.packet_data = self.data_list[5 + self.count_options]
            """TCP数据报数据部分"""

        def get_report(self, showdata=False):
            print('TCP'.center(40, '='))
            print('字节长度\t: ' + str(self.length))
            print('首部'.center(30, '-'))
            print('源端口\t: ' + str(self.src_port))
            print('目的端口\t: ' + str(self.des_port))
            print('序号\t\t: ' + str(self.seq_number))
            print('确认号\t: ' + str(self.ack_number))
            print('数据偏移\t: ' + str(self.data_offset * 4))
            print('紧急数据\t: ' + ('是' if self.urg else '否'))
            print('包含确认\t: ' + ('是' if self.ack else '否'))
            print('尽快交付\t: ' + ('是' if self.psh else '否'))
            print('严重错误\t: ' + ('是' if self.pst else '否'))
            print('建立连接\t: ' + ('是' if self.syn else '否'))
            print('释放连接\t: ' + ('是' if self.fin else '否'))
            print('窗口大小\t: ' + str(self.win_size))
            print('总校验和\t: ' + Convert.to_str(self.checksum, 16))
            print('紧急指针\t: ' + str(self.urg_pointer))
            print('可选字段\t: ' + str(self.options))
            if showdata:
                print('数据'.center(30, '-'))
                hex_dump(self.packet_data)

    def resolve(self):
        """
        解析数据包
        """
        try:
            # 以pcap格式解析数据，
            self.pcap = self._PCAP(self.data)
            # 以ip协议解析数据
            self.ip = self._IP(self.pcap.packet_data)
            # 以tcp协议解析数据
            self.tcp = self._TCP(self.ip.packet_data)
            self.pcap.get_report()
            self.ip.get_report()
            self.tcp.get_report(True)
        except ValueError:
            return False


if __name__ == '__main__':
    hex_dump(Convert.to_byte(0x12345678))
